import whisper
import logging
import os
from opencc import OpenCC

# 设置日志
logging.basicConfig(level=logging.INFO, format='[transcriber] %(asctime)s - %(levelname)s - %(message)s')

def get_project_model_path():
    """获取项目目录下的模型路径"""
    project_dir = os.path.dirname(os.path.abspath(__file__))
    model_dir = os.path.join(project_dir, "whisper_model")
    os.makedirs(model_dir, exist_ok=True)
    return model_dir

# medium    base
def clean_and_reload_model(model_name="medium"):
    """清理损坏的模型并重新加载"""
    # 获取项目目录下的模型路径
    download_root = get_project_model_path()
    model_file = os.path.join(download_root, f"{model_name}.pt")

    # 如果模型文件存在但损坏，则删除它
    if os.path.exists(model_file):
        logging.info(f"删除损坏的模型文件: {model_file}")
        os.remove(model_file)

    # 重新加载模型，指定下载路径
    logging.info(f"重新下载模型到项目目录: {model_name}")
    model = whisper.load_model(model_name, download_root=download_root)
    return model

def transcribe_audio(audio_path, model_name="medium"):
    """使用 Whisper 进行语音识别，返回带有时间戳的文本"""
    # 设置模型下载路径为项目目录
    download_root = get_project_model_path()
    logging.info(f"正在加载 Whisper 模型: {model_name} (路径: {download_root})")

    try:
        model = whisper.load_model(model_name, download_root=download_root)
    except RuntimeError as e:
        if "SHA256 checksum does not match" in str(e):
            logging.warning("检测到模型文件损坏，正在重新下载...")
            model = clean_and_reload_model(model_name)
        else:
            raise e

    logging.info(f"正在识别音频文件: {audio_path}")
    result = model.transcribe(audio_path, language="zh")

    # 创建繁简转换器实例
    cc = OpenCC('t2s')  # 繁体转简体

    # 对每个识别结果进行繁简转换
    for segment in result["segments"]:
        segment["text"] = cc.convert(segment["text"])

    # 拼接所有文字成一段话并在控制台输出
    full_text = "".join([segment["text"] for segment in result["segments"]])
    logging.info(f"完整识别文本: {full_text}")
    logging.info(f"语音识别完成，识别结果示例: {result['segments'][0] if result['segments'] else '无内容'}")
    return result["segments"]
